import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
import os
# targetfolder
import cv2
import multiprocessing as mp
# writer = tf.python_io.TFRecordWriter(targetfolder + "train_" + str(tfrecords_file_num) + ".tfrecord")
# 1. read the ground true file
counts = 100000
def generate_tfrecord():
    # 1. use to generate tfrecords
    def process_tf(start,end):
        writer = tf.io.TFRecordWriter("/home/wujing/dataset/imagenet2012_test_"+str(start)+"_"+str(end)+".tfrecord")
        image_names = []
        for k in range(start, end):
            image_names.append("ILSVRC2012_test_" + f"{k:08}" + ".JPEG")

        for k in range(len(image_names)):
            image = cv2.imread(image_path + "/" + image_names[k])
            record_bytes = tf.train.Example(features=tf.train.Features(feature={
                "image": tf.train.Feature(float_list=tf.train.FloatList(value=image.flatten())),
                "shape": tf.train.Feature(int64_list=tf.train.Int64List(value=image.shape))})).SerializeToString()
            writer.write(record_bytes)

    # 2. read image
    image_path = "/home/wujing/sotfware/dataset/ILSVRC2012_img_test/test"

    # 3. use multiple process to accelerate the process
    mp.Process(target=process_tf, args=(1, 1000)).start()
    mp.Process(target=process_tf, args=(1000, 2000)).start()
    mp.Process(target=process_tf, args=(2000, 3000)).start()
    mp.Process(target=process_tf, args=(3000, 4000)).start()
    mp.Process(target=process_tf, args=(4000, 5000)).start()
    mp.Process(target=process_tf, args=(5000, 6000)).start()
    mp.Process(target=process_tf, args=(6000, 7000)).start()
    mp.Process(target=process_tf, args=(7000, 8000)).start()
    mp.Process(target=process_tf, args=(8000, 9000)).start()
    mp.Process(target=process_tf, args=(9000, 10001)).start()


def read_tfrecords(start,end):

    def _parse_image_function(example_proto):
        image_feature_description = {
            "image": tf.io.FixedLenFeature([], tf.float32),
            "shape": tf.io.FixedLenFeature([], tf.int64)
        }
        # Parse the input tf.train.Example proto using the dictionary above.
        return tf.io.parse_single_example(example_proto, features=image_feature_description)

    tf_path = "/home/wujing/dataset/"
    tf_name = "imagenet2012_test_"+str(start)+"_"+str(end)+".tfrecord"
    image_dataset = tf.data.TFRecordDataset(tf_path+tf_name).range(2)
    for data in image_dataset:
        print(data)
    '''
    dataset_train = image_dataset.map(_parse_image_function)
    parsed_image_dataset = image_dataset.map(_parse_image_function)
    iterator = tf.compat.v1.data.Iterator.from_structure(dataset_train.output_types, dataset_train.output_shapes)
    image_bbox = iterator.get_next()
    
    result = tf.compat.v1.data.get_output_classes(image_dataset)
    print(result)
    
    # it = iter(image_dataset)
    # print(next(it).numpy())
    i = 0
    for element in image_dataset:
        print(type(element["image"]))
    
    print(parsed_image_dataset)
    for image_features in parsed_image_dataset:
        image = image_features["image"].numpy()
        shape = image_features["shape"].numpy()
        print(image.shape,shape)
        break
    '''


def merge_tfrecords2_dataset():
    pass
read_tfrecords(1,1000)








